{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 读写文件"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 加载和保存张量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0, 1, 2, 3])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "x=torch.arange(4)\n",
    "#保存\n",
    "torch.save(x,\"test_saved_x_File\")\n",
    "#读取\n",
    "load_x=torch.load('test_saved_x_File')\n",
    "print(load_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([0, 1, 2, 3]), tensor([0., 0., 0., 0., 0.]))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([0., 1., 2., 3., 0., 0., 0., 0., 0.])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#存储一个张量列表，然后读取回内存\n",
    "y=torch.zeros(5)\n",
    "torch.save([x,y],'x_y_f')#保存\n",
    "#读取\n",
    "x2,y2=torch.load('x_y_f')\n",
    "print((x2,y2))\n",
    "torch.cat([x,y2])#合并俩张量\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'dict'>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'x': tensor([0, 1, 2, 3]), 'y': tensor([0., 0., 0., 0., 0.])}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"保存为字典映射\"\"\"\n",
    "my_dict={\"x\":x,'y':y}\n",
    "torch.save(my_dict,'mydict')\n",
    "load_mydict=torch.load('mydict')\n",
    "print(type(load_mydict))\n",
    "load_mydict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 加载和保存模型参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class mlp(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.hidden=nn.Linear(20,256)\n",
    "        \n",
    "        self.output=nn.Linear(256,10)\n",
    "    def forward(self,x):\n",
    "        return self.output(F.relu(self.hidden(x)))\n",
    "X=torch.rand(2,20)\n",
    "net=mlp()\n",
    "Y=net(X)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "将模型参数存储在一个叫mlp.params的文件中"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(net.state_dict(),'mlp.params')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "读取模型参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.1593, -0.1359, -0.0479, -0.0364, -0.1475, -0.1964,  0.1013, -0.0612,\n",
       "         0.0829, -0.0900, -0.1617,  0.0645, -0.1530,  0.1169,  0.0992, -0.0055,\n",
       "         0.0058, -0.1918,  0.0168,  0.0376], grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clone=mlp()\n",
    "clone.load_state_dict(torch.load('mlp.params'))\n",
    "clone.hidden.weight[0]#隐藏层第一个权重列表（一共256个）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"与原来模型的计算结果相同\"\"\"\n",
    "Y_clone=clone(X)\n",
    "Y_clone==Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sat Jan 18 20:09:57 2025       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 512.72       Driver Version: 512.72       CUDA Version: 11.6     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  NVIDIA GeForce ... WDDM  | 00000000:01:00.0 Off |                  N/A |\n",
      "| N/A   50C    P8     4W /  N/A |      0MiB /  6144MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|    0   N/A  N/A     11924      C   ...rograms\\Ollama\\ollama.exe    N/A      |\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在PyTorch中，CPU和GPU可以用torch.device('cpu') 和torch.device('cuda')表示。应该注意的是，cpu设</br>\n",
    "备意味着所有物理CPU和内存，这意味着PyTorch的计算将尝试使用所有CPU核心。然而，gpu设备只代表一</br>\n",
    "个卡和相应的显存。如果有多个GPU，我们使用torch.device(f'cuda:{i}') 来表示第i块GPU（i从0开始）。</br>\n",
    "另外，cuda:0和cuda是等价的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cpu'), device(type='cuda'), device(type='cuda', index=1))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.device(\"cpu\"),torch.device('cuda'),torch.device('cuda:1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"查询可用GPU数量\"\"\"\n",
    "torch.cuda.device_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"若是有GPU，返回GPU，否则返回CPU\"\"\"\n",
    "def try_gpu(i=0):\n",
    "    if torch.cuda.device_count()>=i+1:\n",
    "        return torch.device(f'cuda:{i+1}')\n",
    "    return torch.device('cpu')\n",
    "\"\"\"返回所有可用的GPU，如果没有GPU返回[cpu(),]\"\"\"\n",
    "def try_all_gpus():\n",
    "    devices=[torch.device(f'cuda:{i}')\n",
    "                for i in range(torch.cuda.device_count())\n",
    "             ]\n",
    "    return devices if devices else [torch.device('cpu')]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cuda', index=1),\n",
       " device(type='cpu'),\n",
       " [device(type='cuda', index=0)])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "try_gpu(),try_gpu(10),try_all_gpus()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### GPU与张量"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通常，张量创建默认在CPU上"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x=torch.tensor([1,2,3])\n",
    "x.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"在GPU上创建张量\"\"\"\n",
    "X=torch.rand(2,3,device='cuda:0')\n",
    "X.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y=torch.rand(2,3,device='cuda:0')\n",
    "Y.device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**神经网络与GPU**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.hidden=nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(20,64),nn.ReLU(),\n",
    "            nn.Linear(64,128),nn.ReLU(),\n",
    "            nn.Linear(128,32),nn.ReLU()\n",
    "        )\n",
    "        self.out=nn.Linear(32,1)\n",
    "    def forward(self,x):\n",
    "        return self.out(self.hidden(x))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用GPU加载数据：实际时间11ms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: total: 15.6 ms\n",
      "Wall time: 11 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[0.1213],\n",
       "        [0.1248],\n",
       "        [0.1246],\n",
       "        ...,\n",
       "        [0.1365],\n",
       "        [0.1345],\n",
       "        [0.1129]], device='cuda:0', grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\"\"\"数据和模型都需要放入gpu\"\"\"\n",
    "X=torch.rand(500000,20,device='cuda:0')\n",
    "NET_MLP=MLP()\n",
    "NET_MLP_gpu=NET_MLP.to(device='cuda:0')\n",
    "NET_MLP_gpu(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用CPU读取数据：实际时间318ms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: total: 1.64 s\n",
      "Wall time: 318 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[0.1450],\n",
       "        [0.1594],\n",
       "        [0.1609],\n",
       "        ...,\n",
       "        [0.1596],\n",
       "        [0.1448],\n",
       "        [0.1692]], grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "X_cpu=torch.rand(500000,20)\n",
    "MLP_CPU=MLP()\n",
    "MLP_CPU(X_cpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**清空缓存**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "# 强制进行垃圾回收\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()  # 清空缓存\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DL_pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
